#include <cuda_runtime.h>


template <const int BM, const int BN, const int BK, const int TM>
__global__ void sgemm_BM_BN_BK_TM(int M, int N, int K, float alpha,
                                  const float *A, const float *B,
                                  float beta, float *C) {
    const uint threadColB = threadIdx.x % BN;
    const uint threadRowB = threadIdx.x / BN;

    const float A_multiplier = blockIdx.x * BM + threadIdx.x / BK < M ? 1. : 0.;
    const float B_multiplier = blockIdx.y * BN + threadColB < N ? 1. : 0.;
    const uint A_idx = A_multiplier > 0 ? threadIdx.x / BK * K + (threadIdx.x % BK) : 0;
    const uint B_idx = B_multiplier > 0 ? threadRowB * N + threadColB : 0;

    __shared__ float As[BM * BK];
    __shared__ float Bs[BK * BN];
    float regC[TM] = {0.};      // for each thread, we compute TM results

    // a block will compute BM * BN results
    A += blockIdx.x * BM * K;     // move to (blockIdx.x * BM, 0)
    B += blockIdx.y * BN;       // move to (0, blockIdx.y * BN) 
    C += (blockIdx.x * BM + threadRowB) * N + blockIdx.y * BN + threadColB; // move to (blockIdx.x * BM + threadRowB, blockIdx.y * BN + threadColB)

    for (int k = 0; k < K / BK; k ++) {
        As[threadIdx.x] = A[A_idx] * A_multiplier;
        Bs[threadIdx.x] = B[B_idx] * B_multiplier;

        __syncthreads();
        A += BK;        // move to (blockIdx.x * BM, k)
        B += BK * N;    // move to (k, blockIdx.y * BN)

        for (uint j = 0; j < BK; j ++) {
            for (uint i = 0; i < TM; i ++) {
                regC[i] += As[(i * (BM / TM) + threadRowB) * BK + j] * Bs[j * BN + threadColB];
            }
        }
        __syncthreads();
    }

    const uint remainder = K % BK;
    if (remainder > 0) {
        As[threadIdx.x] = (threadIdx.x % BK) < remainder ? A[A_idx] * A_multiplier : 0;
        Bs[threadIdx.x] = threadRowB < remainder ? B[B_idx] * B_multiplier : 0;
        __syncthreads();
        for (uint j = 0; j < BK; j ++) {
            for (uint i = 0; i < TM; i ++) {
                regC[i] += As[(i * (BM / TM) + threadRowB) * BK + j] * Bs[j * BN + threadColB];
            }
        }
    }

    // write results
    uint threadRowC = blockIdx.x * BM + threadRowB;
    uint strides = BM / TM;
    for (uint i = 0; i < TM; i ++) {
        if (threadRowC < M && B_multiplier > 0) {
            C[0] = alpha * regC[i] + beta * C[0];
            C += strides * N;
            threadRowC += strides;
        }
    }
    
}